import yaml
import joblib
import pandas as pd
import xgboost as xgb
from sklearn.ensemble import RandomForestClassifier


from . training_random_forest  import train_random_forest


def load_config(config_path):
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def check_winning(ticket, winning_numbers, rules):
    red_ticket, blue_ticket = set(ticket[:-1]), ticket[-1]
    red_winning, blue_winning = set(winning_numbers[:-1]), winning_numbers[-1]
    
    red_matches = len(red_ticket.intersection(red_winning))
    blue_matches = 1 if blue_ticket == blue_winning else 0
    
    for rule in rules['winning_conditions']:
        for condition in rule.get('conditions', [{'red_matches': rule['red_matches'], 'blue_matches': rule['blue_matches']}]):
            if red_matches == condition['red_matches'] and blue_matches == condition['blue_matches']:
                return rule['name']
    
    return "未中奖"


def train_model(config, model_name):
    model_config = config['models'][model_name]
    data = pd.read_csv(config['data']['processed_data_path'])
    X = data.drop(columns=[config['data']['target_column']])
    y = data[config['data']['target_column']]

    if model_name == 'random_forest':
        model = RandomForestClassifier(n_estimators=model_config['n_estimators'],
                                       max_depth=model_config['max_depth'],
                                       random_state=model_config['random_state'])
        model.fit(X, y)
        joblib.dump(model, model_config['model_path'])
        print(f"Random Forest 模型已保存到 {model_config['model_path']}")

    elif model_name == 'xgboost':
        dtrain = xgb.DMatrix(X, label=y)
        params = {
            'max_depth': model_config['max_depth'],
            'eta': model_config['learning_rate'],
            'objective': 'binary:logistic',
            'eval_metric': 'logloss',
            'seed': model_config['random_state']
        }
        model = xgb.train(params, dtrain, num_boost_round=model_config['num_boost_round'])
        joblib.dump(model, model_config['model_path'])
        print(f"XGBoost 模型已保存到 {model_config['model_path']}")

    elif model_name == 'lstm':
        X = X.values.reshape(-1, *model_config['input_shape'])
        model = Sequential()
        model.add(LSTM(model_config['units'], activation=model_config['activation'], input_shape=(X.shape[1], X.shape[2])))
        model.add(Dense(1))
        model.compile(optimizer=model_config['optimizer'], loss=model_config['loss'])
        model.fit(X, y, epochs=model_config['epochs'], batch_size=model_config['batch_size'], verbose=1)
        model.save(model_config['model_path'])
        print(f"LSTM 模型已保存到 {model_config['model_path']}")


def check_winning(config,ticket, winning_numbers):
    result = check_winning(ticket, winning_numbers, config['rules'])
    return result


def train_ssq_model():
    train_random_forest()

def train_zc_model():
    pass

def train_qxc_model():
    pass
    

# if __name__ == '__main__':
    
#     config_path = 'config/ssq_config.yaml'
#     config = load_config(config_path)
    # 训练随机森林模型
    # train_model(config, 'random_forest')

    # 训练XGBoost模型
    # train_model(config, 'xgboost')

    # 训练LSTM模型
    # train_model(config, 'lstm')
    
    # 检查彩票中奖情况
    # ticket = [1, 2, 3, 4, 5, 6, 7]  # 示例票
    # winning_numbers = [1, 2, 3, 4, 5, 6, 7]  # 示例开奖号码
    # result = check_winning(config,ticket, winning_numbers)
    # print(f"结果：{result}")